package com.jxb;

import javassist.*;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;

public class PerformMonitorTransformer implements ClassFileTransformer {

    private static final Set<String> classNameSet = new HashSet<>();
    static {
        classNameSet.add("com.jxb.HelloService");
    }

    @Override
    public byte[] transform(ClassLoader loader,
                            String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) {
        String currentClassName = className.replaceAll("/", ".");
        //只增强set中含有的类进行过滤
        if (!classNameSet.contains(currentClassName)) {
            return null;
        }
        System.out.println("transform: [" + currentClassName + "]");

        CtClass ctClass = null;
        try {
            ctClass = ClassPool.getDefault().get(currentClassName);
        } catch (NotFoundException e) {
            e.printStackTrace();
        }
        CtBehavior[] methods = ctClass.getDeclaredBehaviors();
        for (CtBehavior method : methods) {
            try {
                enhanceMethod(method);
            } catch (ClassNotFoundException e) {
                e.printStackTrace();
            } catch (CannotCompileException e) {
                e.printStackTrace();
            }
        }
        try {
            return ctClass.toBytecode();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (CannotCompileException e) {
            e.printStackTrace();
        }
        return null;
    }

    private void enhanceMethod(CtBehavior method) throws ClassNotFoundException, CannotCompileException {
        if (method.isEmpty()) {
            return;
        }
        String name = method.getName();
        //不增强main方法
        if (name.equalsIgnoreCase("main")) {
            return;
        }
        //增强@Bojiao修饰的方法
        if (null == method.getAnnotation(Bojiao.class)) {
            return;
        }

        final StringBuilder source = new StringBuilder();
        source.append("{")
                // 前置增强: 加入时间戳
                .append("long start = System.currentTimeMillis();\n")
                // 保留原有的代码处理逻辑
                .append("$_ = $proceed($$);\n")
                .append("System.out.print(\"method:[" + name + "]\");")
                .append("\n")
                // 后置增强
                .append("System.out.println(\" cost:[\" +(System.currentTimeMillis() -start)+ \"ms]\");")
                .append("}");
        ExprEditor exprEditor = new ExprEditor() {
            @Override
            public void edit(MethodCall methodCall) throws CannotCompileException {
                methodCall.replace(source.toString());
            }
        };
        method.instrument(exprEditor);
    }
}
